from ultralytics import YOLOWorld
import matplotlib.pyplot as plt
import os
import numpy as np
from PIL import Image
from tqdm import tqdm


yolo_model = YOLOWorld("ft_100.pt")
def yolo_annotate(img):

    custom_classes = ["drain well", "persons wearing safety belt", "persons standing on the ladder", "persons with arm exposed", "persons crossing fence"] # person, hole, arm, fense, belt
    yolo_model.set_classes(custom_classes)

    results = yolo_model.predict(img, conf=0.089, iou=0.1) # iou越小，nms过滤越严格，换句话说，iou是能接受的最大重合度
    im_array = results[0].plot()  # plot a BGR numpy array of predictions
    return im_array


if __name__ == "__main__":

    # img_path = r"C:\Users\aikedaer\Desktop\电科院人工智能比赛\人工智能大赛-预赛数据-安监场景\跨越（下穿）安全围栏_fence"
    img_path = r"C:\Users\aikedaer\Desktop\电科院人工智能比赛\traindataset\images"
    # img_path = r"C:\Users\aikedaer\Desktop\sxgl"

    # sb = os.listdir(img_path)
    # for x in sb:
    #     ig_p = os.path.join(img_path, x)
    #     imgs = os.listdir(ig_p)
    #     imgs = [e for e in imgs if e.endswith(".jpg")]
    #     # gt_imgs = [e for e in imgs if e.endswith("_annotated.jpg")]
    #     # raw_imgs = list(set(imgs) - set(gt_imgs))
    #     for p in tqdm(imgs):
    #         # try: 
    #         p_path = os.path.join(ig_p, p)
    #         # gt_path = p_path.replace(".jpg", "_annotated.jpg")
    #         # image_gt = np.asanyarray(Image.open(gt_path))
    #         label_img = yolo_annotate(p_path)
    #         # plt.imshow(image_gt)
    #         # plt.show()
    #         # plt.savefig(label_img[:,:,::-1], )
    #         plt.imshow(label_img[:,:,::-1])
    #         plt.show()

    imgs = os.listdir(img_path)
    imgs = [e for e in imgs if e.endswith(".jpg")]

    for p in tqdm(imgs):
        p_path = os.path.join(img_path, p)
        label_img = yolo_annotate(p_path)
        plt.imshow(label_img[:,:,::-1])
        plt.show()
